from ctypes import CDLL,byref,create_string_buffer
from ctypes import cdll
import json
import base64
import time
import os.path as osp
import os
import io
from PIL import Image
import numpy as np
import sys
import cv2
from tqdm import tqdm
from tool import filesystem, via_tool, opencv_tool # export PYTHONPATH=$PYTHONPATH:`pwd`

class Rec:
    def init(self):
        optimize_recog_dll_path = '/home/xc/work/code/car_cnt/build/src/libcar_cnt_sdk.so'
        self.optimize_recog_dll = cdll.LoadLibrary(optimize_recog_dll_path)

        model_path = "/home/xc/work/code/car_cnt/model"
        model_path_buffer = create_string_buffer(model_path.encode('utf-8'), 65536)

        out = self.optimize_recog_dll.init_model_py(byref(model_path_buffer))
        print("init_model out: ", out)
        return out

    def car_cnt_py(self, imgData):
        img_dir_buffer0 = create_string_buffer(imgData, 65536*500)
        out_info = create_string_buffer(65536*560)
        out_2 = self.optimize_recog_dll.car_cnt_py(byref(img_dir_buffer0), byref(out_info))
        return out_2, json.loads(out_info.value.decode('utf8'))

rec = Rec()
ret = rec.init()
print("ret: ", ret)

def rec_to_via(frame):
    buffer = io.BytesIO()
    rgb_img = frame[:, : , ::-1]
    pil_image = Image.fromarray(rgb_img) 
    pil_image.save(buffer, format="jpeg")
    encoded = base64.b64encode(buffer.getvalue())
    ret, out_data = rec.car_cnt_py(encoded)
    if ret != 0: return ret, None

    regions = []
    for item in out_data["items"]:
        region = dict()
        shape_attribute = {"name": "rect",
                            "x": item["x"],
                            "y": item["y"],
                            "width": item["w"],
                            "height": item["h"]
                        }
        region["shape_attributes"] = shape_attribute
        region["region_attributes"] = {"label": item["category"]}
        regions.append(region)
    return ret, {"regions": regions, "file_attributes":{}}

def video_to_via(video_dir, save_dir=None, step=3, via_name="via_region_data.json"):
    """
    dest_dir: str
        将视频转成图片 并保存到对应目录下 %未实现%
    step: int
        跳帧
    """
    if osp.isfile(video_dir):
        video_paths = [video_dir]
    else:
        video_paths = filesystem.get_all_filepath(video_dir, [".mp4", ".avi"])
    
    ## save image
    for video_path in video_paths:
        base_name  = os.path.basename(video_path)
        new_dir = osp.join(osp.dirname(video_path), ".".join(base_name.split('.')[:-1]))
        if save_dir:
            new_dir = new_dir.replace(video_dir, save_dir)
        if osp.exists(new_dir):
            print("exist... ", new_dir)
            continue
        if not os.path.exists(new_dir):
            os.makedirs(new_dir, exist_ok=True)

        cap = cv2.VideoCapture(video_path)
        count = 0
        total_data = dict()
        while 1:
            ret, frame = cap.read()
            count += 1
            # cv2.imshow("capture",frame)
            # cv2.moveWindow('capture',100,100)
            if not ret:
                print("cap.read() {} is over...".format(base_name))
                break

            # if count < 45223:continue
            # if cv2.waitKey(1) & 0xFF == ord('c'):
            if count % step == 0:
                ret, data = rec_to_via(frame)
                if ret != 0: continue
                file_name = '{}_{}.jpg'.format(count, base_name)
                save_path = osp.join(new_dir, file_name)
                cv2.imwrite(save_path, frame) # [:, 420:1500, :]
                file_size = osp.getsize(save_path)
                data["filename"] = file_name
                data["size"] = file_size
                total_data[file_name + str(file_size)] = data

            # print(np.mean(frame))
            print(count)
            # if cv2.waitKey(1) & 0xFF == ord('b'):
                # break
        cap.release()
        # cv2.destroyAllWindows()
        with open(new_dir + os.sep + via_name, "w") as wf:
            wf.write(json.dumps(total_data))

def images_to_via(data_dir, via_name="via_region_data.json"):
    """
    data_dir: str
    """
    dir_paths = filesystem.get_last_dir(video_dir)
    if len(dir_paths) == 0:
        dir_paths = [data_dir]

    ## save image
    for dir_path in dir_paths:

        total_data = dict()
        for img_path in filesystem.get_all_filepath(dir_path, [".jpg", ".png"]):
            base_name = osp.basename(img_path)
            cv_img = cv2.imread(img_path)
            if cv_img is None:
                print("None...", img_path)
                continue

            ret, data = rec_to_via(cv_img)
            if ret != 0: continue
            file_size = osp.getsize(img_path)
            data["filename"] = base_name
            data["size"] = file_size
            total_data[base_name + str(file_size)] = data

        with open(dir_path + os.sep + via_name, "w") as wf:
            wf.write(json.dumps(total_data))


if __name__ == "__main__":
    video_dir = "/home/xc/work/data/car/2021-06-21"
    save_dir = "/home/xc/work/data/car/2021-06-21"
    step = 2
    video_to_via(video_dir, save_dir, step)

    # video_dir = "/home/xc/work/code/paddle/train_data/det/car/images/05-25/113"
    # images_to_via(video_dir)

